from ultralytics import YOLO
import os
from vehicle_flow_counting import V_OUT_FILE

class YoloTracker:
    def __init__(self, model_path, device="cuda:0"):
        self.model = YOLO(model_path)
        self.device = device

        self.half = True
        self.imgsz = 1920 # 1920 (1280,704)
        self.stream = True
        self.stream_buffer = True

        self.verbose = True

    def track(self, video_path, classes=None, interval_frame=2):
        """
        使用 YOLO 自带 tracker，对视频逐帧处理
        返回：生成器，每帧返回 r 对象
        """
        results = self.model.track(
            device = self.device, #默认None； "cuda:0"
            half = self.half, # 默认 False 控制是否使用 FP16 半精度浮点数 推理; 影响速度和精度
            imgsz = self.imgsz, # 默认640, 1280~3840, 1920, 2560 ; 影响速度和精度
            stream = self.stream, # 默认stream=False: 一次性处理视频的所有帧; True: 需要在循环里处理每帧结果
            stream_buffer = self.stream_buffer,# False 默认 控制视频流处理时是否缓存所有帧
            verbose = self.verbose, # 显示每帧的推理日志

            source = video_path, # 支持 Video file in formats like MP4, AVI, etc. URL for streaming protocols such as RTSP, RTMP, TCP, or an IP address. https://docs.ultralytics.com/modes/predict/#inference-sources
            classes = classes,
            vid_stride = interval_frame, #默认 1 视频帧跳跃间隔

            conf = 0.5, # 默认 0.25
            iou = 0.5, # 默认0.7; iou=0.9（较高阈值）只有几乎重叠的框才会被去掉, 容易留下更多框

            line_width = 2, # 画面窗口字体线宽
            show_labels = True,
            show_conf = True,
            show = False, # 默认 False 不显示画面窗口
            save = False, # 默认  False 存储结果
            project = os.path.dirname(V_OUT_FILE), # 输出保存的根目录（比如 "runs/obb"）
            name = "track_save", # 输出输出保存的子目录（比如 "track/25s.avi"）

            tracker = "bytetrack.yaml" # botsort.yaml 默认; bytetrack.yaml

            # ,persist = True # 默认 False ; 前后帧关联 （试过没啥作用）
        )

        # Results objects have the following methods: https://docs.ultralytics.com/modes/predict/#working-with-results
        for r in results:
            yield r

    @staticmethod
    def parse_tracks(r):
        """
        输入：YOLO track 输出 r
        返回：list of tuples (track_id, cls_id, (cx,cy))
        """
        tracks = []
        if r.obb is not None and r.obb.id is not None:
            ids = r.obb.id.cpu().numpy().astype(int)
            cls = r.obb.cls.cpu().numpy().astype(int)
            xywhr = r.obb.xywhr.cpu().numpy()  # center_x, center_y, w, h, rotation
            for track_id, c, box in zip(ids, cls, xywhr):
                cx, cy = int(box[0]), int(box[1])
                tracks.append((track_id, c, (cx, cy)))
        return tracks


# names: {0: 'pedestrian',
#         1: 'people',
#         2: 'bicycle',
#         3: 'car',
#         4: 'van',
#         5: 'truck',
#         6: 'tricycle',
#         7: 'awning-tricycle',
#         8: 'bus',
#         9: 'motor'}


# for r in results:
#     # print(r)
#     # print(r.speed)
#     print("is_track: ", r.obb.is_track) # True
#     print("track_id: ", r.obb.id) # track id
#     # tensor([1., 2., 3.])
#     print("cls: ", r.obb.cls)
#     # tensor([3., 3., 3.])
#     print("conf: ", r.obb.conf)
#     # tensor([0.9279, 0.9175, 0.9106])
#     print("xywhr: ", r.obb.xywhr) # center point (xy), width, height, and rotation
#     # tensor([[1.8397e+03, 1.1863e+03, 1.0209e+02, 8.1186e+01, 1.5650e+00],
#     #     [2.0989e+03, 1.5744e+03, 9.5682e+01, 5.9354e+01, 1.5637e+00],
#     #     [1.3568e+03, 8.7568e+02, 5.5952e+01, 9.3395e+01, 1.5684e+00]])
#     print("xyxyxyxy: ", r.obb.xyxyxyxy) # 4 Vertex
#     # tensor([[[ 1.8398e+03,  1.3004e+03],
#     #      [ 1.9187e+03,  1.3001e+03],
#     #      [ 1.9182e+03,  1.1951e+03],
#     #      [ 1.8393e+03,  1.1954e+03]],
#     #     [[ 2.1027e+03,  1.7408e+03],
#     #      [ 2.1563e+03,  1.7405e+03],
#     #      [ 2.1559e+03,  1.6450e+03],
#     #      [ 2.1023e+03,  1.6452e+03]],
#     #     [[ 1.5139e+03,  9.2221e+02],
#     #      [ 1.5140e+03,  8.6350e+02],
#     #      [ 1.4230e+03,  8.6331e+02],
#     #      [ 1.4229e+03,  9.2202e+02]]])
